# This code is part of cqlib.
#
# (C) Copyright 2025 China Telecom Quantum Group, QuantumCTek Co., Ltd.,
# Center for Excellence in Quantum Information and Quantum Physics.
#
# This code is licensed under the Apache License, Version 2.0.
# You may obtain a copy of this license in the LICENSE file in the root
# directory of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""Bloch-sphere visualization for cqlib.

Overview:
  * :func:`plot_bloch_vector` draws a single Bloch sphere and a state vector
    arrow (for one qubit). Input may be a Cartesian/ spherical 3-vector, a
    single-qubit statevector, a density matrix, or a cqlib simulator dict.
  * :func:`plot_bloch_multivector` draws one Bloch sphere per qubit given a
    multi-qubit state (statevector/density matrix) or a list of per-qubit
    Bloch vectors.

Key Features:
  * Pure NumPy + Matplotlib (no QuTiP dependency).
  * Accepts either a statevector (shape ``(2**n,)``) or a density matrix
    (shape ``(2**n, 2**n)``), plus convenience inputs.
  * Robust partial trace via ``np.einsum`` to compute single-qubit reduced
    states.
  * Works in scripts and Jupyter; returns figure/axes for further customization.

Coordinate Conventions:
  * Cartesian vectors are given as ``(x, y, z)``.
  * Spherical vectors are ``(r, theta, phi)`` with ``theta`` the polar angle
    measured from +Z, and ``phi`` the azimuth from +X toward +Y.

Examples:
  >>> import numpy as np
  >>> from cqlib.visualization.bloch import plot_bloch_vector, plot_bloch_multivector
  >>> # 1) A single-qubit +X state |+> = (|0>+|1>)/√2
  >>> v = np.array([1.0, 0.0, 0.0])
  >>> fig1, ax = plot_bloch_vector(v, title="|+> on Bloch sphere")
  >>> # 2) Two-qubit entangled state – local Bloch vectors collapse toward 0
  >>> psi = np.zeros(4, complex); psi[0] = psi[3] = 1/np.sqrt(2)
  >>> fig2, axes = plot_bloch_multivector(psi, title="Bell state local Bloch vectors")
  >>> fig1.savefig('bloch_single.png', dpi=200)
  >>> fig2.savefig('bloch_multi.png', dpi=200)

Notes:
  * For highly entangled multi-qubit states, *local* Bloch vectors can be close
    to zero; this indicates mixedness of the subsystem.
  * If you have a ``cqlib`` Circuit, first obtain a statevector or density
    matrix using your preferred simulator, then pass it to these functions.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, List, Optional, Sequence, Tuple, Union
import math
import numpy as np

# Matplotlib is imported lazily inside helpers to avoid import cost for users
# who do not plot. This also keeps cqlib importable in headless environments
# where Matplotlib might be optional.

# -------------------------------------------
# Visual style
# -------------------------------------------

_DEF_STYLE = {
    # Sphere
    "sphere_alpha": 0.12,  # historical, currently unused; kept for compatibility
    "sphere_face_alpha": 0.10,
    "sphere_edgecolor": (0.7, 0.7, 0.7, 0.5),
    # Reference lines
    "equator_color": "#000000",
    "equator_linestyle": "-",
    "equator_linewidth": 1.0,
    "meridian_color": "#000000",
    "meridian_linestyle": "--",
    "meridian_linewidth": 1.0,
    # Axes
    "axes_color": "#444444",
    "axes_linewidth": 1.0,
    "label_color": "#000000",
    # Vector (arrow)
    "vector_color": "#D62728",
    "vector_linewidth": 2.0,
    "vector_arrow_ratio": 0.12,
    # Box decorations
    "hide_box": True,
    "hide_axis_lines": True,
}

# -------------------------------------------
# Public API
# -------------------------------------------

def plot_bloch_vector(
    vec_or_state,
    *,
    ax=None,
    title: Optional[str] = None,
    radius: float = 1.5,
    draw_quiver: bool = True,
    show_axes_labels: bool = True,
    style: Optional[dict] = None,
    coord: str = "auto",  # "auto"/"cart"/"spherical"
    degrees: bool = False,  # spherical coordinates in degrees
):
    """Draw a single Bloch sphere and state vector for one qubit.

    This function accepts either a 3-vector or a single-qubit quantum state.
    When ``coord="spherical"``, a 3-vector is interpreted as ``(r, theta, phi)``
    with the physics convention (theta from +Z, phi from +X toward +Y).

    Args:
        vec_or_state: One of the following inputs:
            - Cartesian Bloch vector ``(x, y, z)`` (default interpretation for a
              length-3 real vector).
            - Spherical vector ``(r, theta, phi)`` when ``coord='spherical'``; set
              ``degrees=True`` to interpret angles in degrees.
            - Single-qubit statevector of shape ``(2,)`` or mapping
              ``{"0"|"1": amplitude}`` returned by the cqlib statevector simulator.
            - Single-qubit density matrix of shape ``(2, 2)``.
        ax (Axes3D | None): Matplotlib 3D axes to draw on. When ``None``, a new
            figure and axis are created.
        title (str | None): Optional plot title.
        radius (float): Radius of the Bloch sphere (visual scale). Axes limits are
            set to ``[-1, 1]``; values >1 visually enlarge the sphere relative to
            the axes frame.
        draw_quiver (bool): Whether to draw the state vector arrow.
        show_axes_labels (bool): Whether to label the X/Y/Z axes.
        style (dict | None): Style overrides; keys are merged with ``_DEF_STYLE``.
        coord (str): Coordinate interpretation for 3-vectors: ``'auto'``/``'cart'`` or
            ``'spherical'``.
        degrees (bool): If ``True`` and ``coord='spherical'``, interpret ``theta`` and
            ``phi`` in degrees.

    Returns:
        tuple[Figure, Axes3D]: The Matplotlib figure and 3D axis objects.

    Raises:
        ValueError: If a multi-qubit state is provided, or if the input shape is
            invalid for any of the accepted formats.

    Examples:
        >>> plot_bloch_vector([0.5, 0.866, 0.0])  # Cartesian
        >>> plot_bloch_vector((1.0, 90, 60), coord="spherical", degrees=True)
    """
    fig, ax = _ensure_3d_axis(ax)
    _draw_sphere(ax, radius=radius, style=style or _DEF_STYLE)
    _draw_reference_axes(ax, radius=radius, show_labels=show_axes_labels, style=style or _DEF_STYLE)
    _apply_axis_box_style(ax, style or _DEF_STYLE)

    # --- detect input kind (vector/state; Cartesian vs spherical)
    v: Optional[np.ndarray] = None

    if isinstance(vec_or_state, dict):
        arr = statevector_dict_to_array(vec_or_state)
        n = _infer_num_qubits(arr.size)
        if n != 1:
            raise ValueError("plot_bloch_vector: multi-qubit state given; use plot_bloch_multivector().")
        rho = np.outer(arr, np.conjugate(arr))
        _, [v] = _state_to_local_bloch_vectors(rho)

    else:
        arr = np.asarray(vec_or_state)
        # Case B: explicit length-3 vector
        if arr.ndim == 1 and arr.shape == (3,) and np.isrealobj(arr):
            mode = (coord or "auto").lower()
            if mode in ("spherical", "sph"):
                v = _spherical_to_cart(arr, degrees=degrees)
            else:  # "auto"/"cart" → Cartesian
                v = _coerce_vec3(arr)
        # Case C: statevector (2,) or density matrix (2,2)
        else:
            if arr.ndim == 1:
                n = _infer_num_qubits(arr.size)
                if n != 1:
                    raise ValueError("plot_bloch_vector: multi-qubit state given; use plot_bloch_multivector().")
                rho = np.outer(arr, np.conjugate(arr))
                _, [v] = _state_to_local_bloch_vectors(rho)
            elif arr.ndim == 2 and arr.shape[0] == arr.shape[1]:
                n = _infer_num_qubits(arr.shape[0])
                if n != 1:
                    raise ValueError("plot_bloch_vector: multi-qubit density matrix; use plot_bloch_multivector().")
                _, [v] = _state_to_local_bloch_vectors(arr)
            else:
                raise ValueError("plot_bloch_vector expects (x,y,z)/(r,θ,φ) or single-qubit statevector/density matrix.")

    # Clamp & draw
    v_norm = np.linalg.norm(v)
    if v_norm > 1.0 and v_norm > 0:
        v = v / v_norm
    if draw_quiver:
        _draw_vector(ax, v, radius=radius, style=style or _DEF_STYLE)

    ax.view_init(elev=22, azim=-60)
    if title:
        ax.set_title(title)
    return fig, ax


def _apply_axis_box_style(ax, style: dict):
    """Apply axes styling by hiding panes and/or box edge lines.

    This function hides the external 3D panes (aka "box faces") and the three
    corner edge lines if requested by ``style``.

    Args:
        ax (Axes3D): Matplotlib 3D axes to modify.
        style (dict): Style dictionary. Recognized keys include:
            - ``hide_box``: If ``True``, hide filled panes and their edges (also
              disables grid lines).
            - ``hide_axis_lines``: If ``True``, hide the three black corner lines
              of the 3D box.
    """
    if style.get("hide_box", True):
        try:
            for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
                pane = getattr(axis, "pane", None)
                if pane is not None:
                    pane.fill = False
                    pane.set_edgecolor((1, 1, 1, 0))
            ax.grid(False)
        except Exception:
            try:
                ax.w_xaxis.set_pane_color((1, 1, 1, 0))
                ax.w_yaxis.set_pane_color((1, 1, 1, 0))
                ax.w_zaxis.set_pane_color((1, 1, 1, 0))
                ax.grid(False)
            except Exception:
                pass

    if style.get("hide_axis_lines", True):
        # Try modern API first
        try:
            ax.xaxis.line.set_color((1, 1, 1, 0))
            ax.yaxis.line.set_color((1, 1, 1, 0))
            ax.zaxis.line.set_color((1, 1, 1, 0))
            ax.xaxis.line.set_linewidth(0.0)
            ax.yaxis.line.set_linewidth(0.0)
            ax.zaxis.line.set_linewidth(0.0)
        except Exception:
            # Fallback to legacy mplot3d attributes
            try:
                ax.w_xaxis.line.set_color((1, 1, 1, 0))
                ax.w_yaxis.line.set_color((1, 1, 1, 0))
                ax.w_zaxis.line.set_color((1, 1, 1, 0))
                ax.w_xaxis.line.set_linewidth(0.0)
                ax.w_yaxis.line.set_linewidth(0.0)
                ax.w_zaxis.line.set_linewidth(0.0)
            except Exception:
                pass


def plot_bloch_multivector(
    state_or_vectors: Union[np.ndarray, Sequence[Sequence[float]], Sequence[complex]],
    *,
    axes: Optional[Sequence] = None,
    figsize: Optional[Tuple[float, float]] = None,
    title: Optional[str] = None,
    labels: Optional[Sequence[str]] = None,
    max_cols: int = 4,
    style: Optional[dict] = None,
    coord: str = "auto",            # "auto"/"cart"/"spherical" when passing vectors
    degrees: bool = False,
) -> Tuple[object, List[object]]:
    """Draw one Bloch sphere per qubit for a state or a list of vectors.

    ``state_or_vectors`` may be a multi-qubit state (statevector/density matrix)
    or an explicit array of vectors shaped ``(n, 3)``.

    Args:
        state_or_vectors: One of the following:
            - Statevector shaped ``(2**n,)`` or mapping ``{"...": amplitude}``.
            - Density matrix shaped ``(2**n, 2**n)``.
            - Array-like of per-qubit vectors shaped ``(n, 3)``. For vectors,
              set ``coord='spherical'`` to interpret rows as ``(r, theta, phi)``.
        axes (Sequence[Axes3D] | None): Iterable of Matplotlib 3D axes with
            length ≥ number of qubits. When ``None``, a new figure with a grid of
            subplots is created.
        figsize (tuple[float, float] | None): Figure size when creating a new figure.
        title (str | None): Global title for the figure.
        labels (Sequence[str] | None): Per-qubit labels; defaults to
            ``['q0', 'q1', ...]`` (reversed so higher-index qubits appear first).
        max_cols (int): Maximum number of columns in the subplot grid.
        style (dict | None): Style overrides forwarded to :func:`plot_bloch_vector`.
        coord (str): Interpretation of vector rows when vectors are provided:
            ``'auto'``/``'cart'`` or ``'spherical'``.
        degrees (bool): If ``True`` and ``coord='spherical'``, interpret angles in degrees.

    Returns:
        tuple[Figure, list[Axes3D]]: The Matplotlib figure and the list of axes.

    Raises:
        ValueError: If the input does not match any accepted shape/format.
    """

    # --- normalize inputs: state vs list of vectors
    if isinstance(state_or_vectors, dict):
        arr = statevector_dict_to_array(state_or_vectors)
        num_qubits, r_list = _state_to_local_bloch_vectors(arr)
    else:
        arr = np.asarray(state_or_vectors)
        # Case A: explicit vectors (n,3)
        if arr.ndim == 2 and arr.shape[1] == 3 and np.isrealobj(arr):
            num_qubits = arr.shape[0]
            if coord.lower() in ("spherical", "sph"):
                r_list = [_spherical_to_cart(row, degrees=degrees) for row in arr]
            else:
                r_list = [_coerce_vec3(row) for row in arr]
        # Case B: statevector (2**n,) or density matrix (2**n,2**n)
        else:
            if arr.ndim == 1:
                num_qubits, r_list = _state_to_local_bloch_vectors(arr)
            elif arr.ndim == 2 and arr.shape[0] == arr.shape[1]:
                num_qubits, r_list = _state_to_local_bloch_vectors(arr)
            else:
                raise ValueError(
                    "plot_bloch_multivector expects statevector/density-matrix or an (n,3) array of vectors."
                )

    if labels is None:
        labels = [f"q{i}" for i in range(num_qubits)][::-1]

    if axes is None:
        import matplotlib.pyplot as plt
        cols = min(max_cols, num_qubits)
        rows = math.ceil(num_qubits / cols)
        fig = plt.figure(figsize=figsize or (3.0 * cols, 3.2 * rows))
        axes = []
        for i in range(num_qubits):
            ax = fig.add_subplot(rows, cols, i + 1, projection="3d")
            axes.append(ax)
    else:
        fig = axes[0].get_figure() if len(axes) else None

    for i, (ax, rvec) in enumerate(zip(axes, r_list)):
        plot_bloch_vector(rvec, ax=ax, title=labels[i], style=style)

    if title and fig is not None:
        fig.suptitle(title)
    return fig, list(axes)


# -------------------------------------------
# Internals – Bloch vectors from (reduced) states
# -------------------------------------------

_SIGMA_X = np.array([[0, 1], [1, 0]], dtype=complex)
_SIGMA_Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
_SIGMA_Z = np.array([[1, 0], [0, -1]], dtype=complex)
_PAULIS = (_SIGMA_X, _SIGMA_Y, _SIGMA_Z)


def _state_to_local_bloch_vectors(state: Union[np.ndarray, Sequence[complex]]):
    """Compute per-qubit Bloch vectors from a state.

    Args:
        state: Statevector of shape ``(2**n,)`` or density matrix of shape
            ``(2**n, 2**n)``.

    Returns:
        tuple[int, list[np.ndarray]]: The number of qubits and the list of
        length-3 real vectors.

    Raises:
        ValueError: If the input dimensionality is not a power of two or has
            the wrong rank/shape.
    """
    arr = np.asarray(state)
    if arr.ndim == 1:
        n = _infer_num_qubits(arr.size)
        rho = np.outer(arr, np.conjugate(arr))  # |psi><psi|
    elif arr.ndim == 2 and arr.shape[0] == arr.shape[1]:
        n = _infer_num_qubits(arr.shape[0])
        rho = arr
    else:
        raise ValueError(
            "state must be a statevector (1D, length 2**n) or a density matrix (2D, 2**n x 2**n)."
        )

    r_list: List[np.ndarray] = []
    for qubit in range(n):
        rho_k = _partial_trace_qubit(rho, keep=qubit, num_qubits=n)
        rvec = np.array([np.real(np.trace(rho_k @ P)) for P in _PAULIS], dtype=float)
        # numerical clean-up for tiny imaginary/epsilon noise
        rvec[np.abs(rvec) < 1e-12] = 0.0
        r_list.append(rvec)
    return n, r_list


def _infer_num_qubits(dim: int) -> int:
    """Infer the number of qubits from a Hilbert-space dimension.

    Args:
        dim (int): Hilbert-space dimension (must be a power of two).

    Returns:
        int: The number of qubits ``n`` such that ``2**n == dim``.

    Raises:
        ValueError: If ``dim`` is not a power of two.
    """
    n = int(round(math.log2(dim)))
    if 2 ** n != dim:
        raise ValueError(f"Dimension {dim} is not a power of two.")
    return n


def _partial_trace_qubit(rho: np.ndarray, *, keep: int, num_qubits: int) -> np.ndarray:
    """Return the 2x2 reduced density matrix of one qubit via partial trace.

    Uses a single ``np.einsum`` call to trace out all but the selected qubit,
    which is robust to axis shrinking (unlike repeated ``np.trace`` calls).

    Args:
        rho (np.ndarray): Density matrix of shape ``(2**n, 2**n)``.
        keep (int): Qubit index to keep. Indexing follows **left-to-right** bit
            order of computational-basis strings; 0 refers to the most
            significant/leftmost bit in a string like ``"q0 q1 ..."``.
        num_qubits (int): Total qubit count ``n``.

    Returns:
        np.ndarray: The 2x2 reduced density matrix of the kept qubit.

    Raises:
        ValueError: If ``rho`` has the wrong shape for ``num_qubits``.
    """
    import string

    dim = 2 ** num_qubits
    if rho.shape != (dim, dim):
        raise ValueError("rho has wrong shape for the provided num_qubits")

    # Reshape into 2n indices: r0..r{n-1} c0..c{n-1}
    rho_t = rho.reshape((2,) * (2 * num_qubits))

    # Build Einstein summation indices. We'll set r_q == c_q for all q != keep
    # so those are traced out. The kept qubit uses distinct labels 'k' and 'K'
    # so the output retains a 2x2 matrix.
    lowers = list(string.ascii_lowercase)
    uppers = list(string.ascii_uppercase)
    if num_qubits > len(lowers):
        raise ValueError("Too many qubits for simple einsum label set.")

    rlabels = lowers[:num_qubits]
    clabels = uppers[:num_qubits]

    for q in range(num_qubits):
        if q == keep:
            rlabels[q] = 'k'
            clabels[q] = 'K'
        else:
            # pair the labels to perform a trace over this qubit
            clabels[q] = rlabels[q]

    subscripts = ''.join(rlabels + clabels)
    out_sub = 'kK'

    rho_k = np.einsum(f'{subscripts}->{out_sub}', rho_t)
    return rho_k


# -------------------------------------------
# Internals – drawing primitives
# -------------------------------------------

def _spherical_to_cart(vec, *, degrees: bool = False) -> np.ndarray:
    """Convert spherical ``(r, theta, phi)`` to Cartesian ``(x, y, z)``.

    Args:
        vec: Tuple/array ``(r, theta, phi)``.
        degrees (bool): If ``True``, interpret angles in degrees; otherwise radians.

    Returns:
        np.ndarray: Length-3 real vector ``(x, y, z)``.
    """
    r, theta, phi = vec
    if degrees:
        theta = np.deg2rad(theta)
        phi = np.deg2rad(phi)
    x = r * np.sin(theta) * np.cos(phi)
    y = r * np.sin(theta) * np.sin(phi)
    z = r * np.cos(theta)
    return np.array([x, y, z], dtype=float)


def _ensure_3d_axis(ax=None):
    """Return a (figure, 3D axis) pair, creating them when needed.

    Args:
        ax (Axes3D | None): Existing 3D axis. When ``None``, a new figure and axis
            are created.

    Returns:
        tuple[Figure, Axes3D]: The Matplotlib figure and 3D axes.
    """
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 – ensure 3D registered

    if ax is None:
        fig = plt.figure(figsize=(3.2, 3.2))
        ax = fig.add_subplot(111, projection="3d")
    else:
        fig = ax.get_figure()
    return fig, ax


def _draw_sphere(ax, *, radius: float = 0.98, meridians: int = 12, parallels: int = 12, style: dict = _DEF_STYLE):
    """Draw the Bloch sphere surface and grid lines.

    Args:
        ax (Axes3D): Matplotlib 3D axes object.
        radius (float): Sphere radius (visual scale within the axes limits).
        meridians (int): Number of meridian samples for the mesh.
        parallels (int): Number of parallel samples for the mesh.
        style (dict): Style dictionary; uses ``sphere_face_alpha`` and ``sphere_edgecolor``.
    """
    u = np.linspace(0, 2 * np.pi, meridians)
    v = np.linspace(0, np.pi, parallels)
    x = radius * np.outer(np.cos(u), np.sin(v))
    y = radius * np.outer(np.sin(u), np.sin(v))
    z = radius * np.outer(np.ones_like(u), np.cos(v))

    ax.plot_surface(x, y, z, rstride=1, cstride=1, linewidth=0.4,
                    edgecolor=style["sphere_edgecolor"], alpha=style["sphere_face_alpha"], shade=True)


def _draw_reference_axes(ax, *, radius: float = 0.98, show_labels: bool = True, style: dict = _DEF_STYLE):
    """Draw reference axes, equator, and a prime meridian.

    Args:
        ax (Axes3D): Matplotlib 3D axes object.
        radius (float): Visual scale used for lines and label placement.
        show_labels (bool): Whether to place X/Y/Z labels at the sphere boundary.
        style (dict): Style dictionary; uses ``axes_*``, ``equator_*``,
            ``meridian_*``, and ``label_color``.
    """
    _line(ax, (-radius, 0, 0), (radius, 0, 0), color=style["axes_color"], linewidth=style["axes_linewidth"])
    _line(ax, (0, -radius, 0), (0, radius, 0), color=style["axes_color"], linewidth=style["axes_linewidth"])
    _line(ax, (0, 0, -radius), (0, 0, radius), color=style["axes_color"], linewidth=style["axes_linewidth"])

    t = np.linspace(0, 2 * np.pi, 200)
    ax.plot(radius * np.cos(t), radius * np.sin(t), 0 * t,
            linewidth=style["equator_linewidth"], linestyle=style["equator_linestyle"], color=style["equator_color"])
    ax.plot(radius * 0 * t, radius * np.sin(t), radius * np.cos(t),
            linewidth=style["meridian_linewidth"], linestyle=style["meridian_linestyle"], color=style["meridian_color"])

    ax.set_box_aspect((1, 1, 1))
    ax.set_xlim([-1, 1]); ax.set_ylim([-1, 1]); ax.set_zlim([-1, 1])
    ax.set_xticks([]); ax.set_yticks([]); ax.set_zticks([])
    if show_labels:
        ax.text(1.05 * radius, 0, 0, "X", color=style["label_color"]) 
        ax.text(0, 1.05 * radius, 0, "Y", color=style["label_color"]) 
        ax.text(0, 0, 1.05 * radius, "Z", color=style["label_color"]) 


def _draw_vector(ax, v: np.ndarray, *, radius: float = 0.98, style: dict = _DEF_STYLE):
    """Draw the state vector arrow.

    Args:
        ax (Axes3D): Matplotlib 3D axes object.
        v (np.ndarray): Length-3 real vector.
        radius (float): Visual scale used to size the arrow.
        style (dict): Style dictionary; uses ``vector_*``.
    """
    ax.quiver(0, 0, 0,
              float(v[0]) * radius, float(v[1]) * radius, float(v[2]) * radius,
              arrow_length_ratio=style["vector_arrow_ratio"], linewidth=style["vector_linewidth"],
              color=style["vector_color"])


def _line(ax, p0, p1, **kwargs):
    """Draw a line segment between two 3D points.

    Args:
        ax (Axes3D): Matplotlib 3D axes object.
        p0: Start point ``(x, y, z)``.
        p1: End point ``(x, y, z)``.
        **kwargs: Additional keyword arguments forwarded to ``Axes3D.plot``.
    """
    ax.plot([p0[0], p1[0]], [p0[1], p1[1]], [p0[2], p1[2]], **kwargs)


def _coerce_vec3(vec: Sequence[float]) -> np.ndarray:
    """Validate and convert an input sequence to a length-3 float array.

    Args:
        vec: Iterable of three real numbers.

    Returns:
        np.ndarray: Length-3 float array ``(x, y, z)``.

    Raises:
        ValueError: If the input does not have length 3.
    """
    arr = np.asarray(list(vec), dtype=float)
    if arr.shape != (3,):
        raise ValueError("vec must be a length-3 iterable (x, y, z)")
    return arr


# -------------------------------------------
# Compatibility helpers (optional)
# -------------------------------------------

@dataclass
class BlochVectors:
    """Container for per-qubit Bloch vectors.

    Attributes:
        num_qubits (int): Total number of qubits.
        vectors (list[np.ndarray]): Length-3 real arrays, one per qubit.
    """

    num_qubits: int
    vectors: List[np.ndarray]


def state_to_bloch_vectors(state: Union[np.ndarray, Sequence[complex]]) -> BlochVectors:
    """Return per-qubit Bloch vectors computed from a state (no plotting).

    Args:
        state: Statevector shaped ``(2**n,)`` or density matrix shaped
            ``(2**n, 2**n)``.

    Returns:
        BlochVectors: Dataclass with the number of qubits and the per-qubit vectors.
    """
    n, rlist = _state_to_local_bloch_vectors(state)
    return BlochVectors(num_qubits=n, vectors=rlist)


def statevector_dict_to_array(state_dict: dict[str, complex]) -> np.ndarray:
    """Convert ``StatevectorSimulator.statevector()`` dict to a 1D array.

    The simulator returns a mapping from bitstrings (e.g., ``"010"``) to complex
    amplitudes. This helper orders the entries by the integer value of the
    bitstring (i.e., ``int(bitstring, 2)``), producing an array compatible with
    the plotting functions in this module.

    Args:
        state_dict: Mapping like ``{"000": amp0, "001": amp1, ...}``.

    Returns:
        np.ndarray: Complex array of shape ``(2**n,)``.
    """
    if not state_dict:
        return np.asarray([], dtype=complex)
    n = len(next(iter(state_dict.keys())))
    size = 2 ** n
    arr = np.zeros(size, dtype=complex)
    for b, amp in state_dict.items():
        arr[int(b, 2)] = amp
    return arr
